# -*- coding: utf-8 -*-
"""
数据库连接管理
"""

import logging
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.exc import SQLAlchemyError
from contextlib import contextmanager
from typing import Generator

from config.settings import DATABASE_CONFIG
from .models import Base

logger = logging.getLogger(__name__)


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, db_type: str = 'sqlite'):
        """
        初始化数据库管理器
        
        Args:
            db_type: 数据库类型，默认为sqlite
        """
        self.db_type = db_type
        self.config = DATABASE_CONFIG[db_type]
        self.engine = None
        self.SessionLocal = None
        self._initialize()
    
    def _initialize(self):
        """初始化数据库连接"""
        try:
            self.engine = create_engine(
                self.config['url'],
                echo=self.config.get('echo', False)
            )
            self.SessionLocal = sessionmaker(
                autocommit=False,
                autoflush=False,
                bind=self.engine
            )
            logger.info(f"数据库连接初始化成功: {self.config['url']}")
        except Exception as e:
            logger.error(f"数据库连接初始化失败: {e}")
            raise
    
    def create_tables(self):
        """创建所有表"""
        try:
            Base.metadata.create_all(bind=self.engine)
            logger.info("数据库表创建成功")
        except Exception as e:
            logger.error(f"数据库表创建失败: {e}")
            raise
    
    def drop_tables(self):
        """删除所有表"""
        try:
            Base.metadata.drop_all(bind=self.engine)
            logger.info("数据库表删除成功")
        except Exception as e:
            logger.error(f"数据库表删除失败: {e}")
            raise
    
    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """
        获取数据库会话的上下文管理器
        
        Yields:
            Session: 数据库会话对象
        """
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except SQLAlchemyError as e:
            session.rollback()
            logger.error(f"数据库操作失败: {e}")
            raise
        except Exception as e:
            session.rollback()
            logger.error(f"未知错误: {e}")
            raise
        finally:
            session.close()
    
    def get_session_direct(self) -> Session:
        """
        直接获取数据库会话（需要手动管理）
        
        Returns:
            Session: 数据库会话对象
        """
        return self.SessionLocal()
    
    def close(self):
        """关闭数据库连接"""
        if self.engine:
            self.engine.dispose()
            logger.info("数据库连接已关闭")


# 全局数据库管理器实例
db_manager = DatabaseManager()


def get_db_session() -> Generator[Session, None, None]:
    """
    获取数据库会话的依赖注入函数
    
    Yields:
        Session: 数据库会话对象
    """
    with db_manager.get_session() as session:
        yield session


def init_database():
    """初始化数据库"""
    try:
        # 创建表
        db_manager.create_tables()
        
        # 插入基础配置数据
        from .init_data import insert_initial_data
        insert_initial_data()
        
        logger.info("数据库初始化完成")
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        raise


if __name__ == "__main__":
    # 设置日志
    logging.basicConfig(level=logging.INFO)
    
    # 初始化数据库
    init_database()
